import chromadb
from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.core import StorageContext
from llama_index.core import Settings
from llama_index.core import SimpleDirectoryReader
from llama_index.core import SummaryIndex, VectorStoreIndex
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.query_engine.router_query_engine import RouterQueryEngine
from llama_index.core.selectors import LLMSingleSelector
from llama_index.core.tools import QueryEngineTool
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.llms.openai import OpenAI

# Create a Chroma client and collection
chroma_client = chromadb.PersistentClient(path="./chroma_db")  # 指定存储路径
try:
    chroma_collection = chroma_client.get_collection("example_collection")
except Exception:
    print(" Collection example_collection does not exist.")
    chroma_collection = chroma_client.create_collection("example_collection")
# Set up the ChromaVectorStore and StorageContext
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)

# 创建 StorageContext
storage_context = StorageContext.from_defaults(vector_store=vector_store)

# 加载已保存的索引
vector_index = VectorStoreIndex.from_vector_store(vector_store, storage_context=storage_context)

# load documents
documents = SimpleDirectoryReader(input_files=["metagpt.pdf"]).load_data()
print(documents)
splitter = SentenceSplitter(chunk_size=1024)
nodes = splitter.get_nodes_from_documents(documents)
print(nodes)
Settings.llm = OpenAI(model="gpt-3.5-turbo")
Settings.embed_model = OpenAIEmbedding(model="text-embedding-ada-002")
summary_index = SummaryIndex(nodes)
vector_index = VectorStoreIndex(nodes, vector_store=vector_store)
print(vector_index)
summary_query_engine = summary_index.as_query_engine(
    response_mode="tree_summarize",
    use_async=True,
)
vector_query_engine = vector_index.as_query_engine()

summary_tool = QueryEngineTool.from_defaults(
    query_engine=summary_query_engine,
    description=(
        "Useful for summarization questions related to MetaGPT"
    ),
)

vector_tool = QueryEngineTool.from_defaults(
    query_engine=vector_query_engine,
    description=(
        "Useful for retrieving specific context from the MetaGPT paper."
    ),
)

query_engine = RouterQueryEngine(
    selector=LLMSingleSelector.from_defaults(),
    query_engine_tools=[
        summary_tool,
        vector_tool,
    ],
    verbose=True
)

response = query_engine.query("What is the summary of the document?")
print(str(response))

print(len(response.source_nodes))

response = query_engine.query(
    "How do agents share information with other agents?"
)
print(str(response))
